from prompt_generators.prompt_registry import register_prompt, BasePromptStrategy
from prompt_generators.prompt_utils import read_relavant_files, generate_template


@register_prompt("tilelang", "add_shot")
class TilelangDefaultPromptStrategy(BasePromptStrategy):        
    def generate(self, op):
        arc_src, example_arch_src, example_new_arch_src = read_relavant_files('tilelang', op, 'add')
        return generate_template(arc_src, example_arch_src, example_new_arch_src, 'tilelang')

